import math
import numpy as np
from collections import Counter
from operator import itemgetter


class KNN:
    def __init__(self, k=3, task='classification'):
        """
        初始化 KNN 模型

        参数:
            k: 近邻数量
            task: 任务类型，'classification' 或 'regression'
        """
        self.k = k
        self.task = task
        self.X_train = None
        self.y_train = None

    def fit(self, X, y):
        """
        训练模型，存储训练数据

        参数:
            X: 训练特征，形状为 (n_samples, n_features)
            y: 训练标签，形状为 (n_samples,)
        """
        self.X_train = X
        self.y_train = y

    def _euclidean_distance(self, a, b):
        """
        计算两个向量之间的欧氏距离
        """
        return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

    def _get_k_neighbors(self, x):
        """
        获取距离目标样本最近的 k 个邻居

        返回:
            neighbors: 最近的 k 个邻居的索引和距离
        """
        distances = []

        # 计算与所有训练样本的距离
        for i, train_sample in enumerate(self.X_train):
            dist = self._euclidean_distance(x, train_sample)
            distances.append((i, dist))

        # 按距离排序并取前 k 个
        distances.sort(key=itemgetter(1))
        neighbors = distances[:self.k]

        return neighbors

    def predict(self, X):
        """
        预测新样本的标签

        参数:
            X: 测试特征，形状为 (n_samples, n_features)

        返回:
            predictions: 预测结果
        """
        predictions = []

        for sample in X:
            # 获取 k 个最近邻居
            neighbors = self._get_k_neighbors(sample)

            if self.task == 'classification':
                # 分类任务：多数投票
                neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
                most_common = Counter(neighbor_labels).most_common(1)
                predictions.append(most_common[0][0])

            elif self.task == 'regression':
                # 回归任务：平均值
                neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
                predictions.append(sum(neighbor_labels) / len(neighbor_labels))

        return predictions

    def predict_proba(self, X):
        """
        分类任务中返回每个类别的概率（仅适用于分类任务）

        参数:
            X: 测试特征，形状为 (n_samples, n_features)

        返回:
            probabilities: 每个样本属于每个类别的概率
        """
        if self.task != 'classification':
            raise ValueError("predict_proba 仅适用于分类任务")

        probabilities = []

        for sample in X:
            # 获取 k 个最近邻居
            neighbors = self._get_k_neighbors(sample)

            # 统计每个类别的数量
            neighbor_labels = [self.y_train[idx] for idx, _ in neighbors]
            label_counts = Counter(neighbor_labels)

            # 计算概率
            prob_dict = {}
            total = len(neighbors)
            for label in set(self.y_train):
                prob_dict[label] = label_counts.get(label, 0) / total

            probabilities.append(prob_dict)

        return probabilities
